#!/usr/bin/env python3
"""
QP-MPC with automatic left/right/top corridor selection for multiple obstacles.

保存文件为: mpc_quad_cvxpy_multi_obstacle.py
依赖: pip install numpy pandas matplotlib cvxpy osqp

说明:
 - 每步枚举障碍物绕行侧(left/right/top)的候选组合（上限限制）。
 - 对每个组合把对应线性不等式加入 QP（保持凸性），求解并记录 cost (prob.value)。
 - 在可行解中选取 cost 最小者并应用其首个控制。
 - 若所有组合不可行，尝试放宽约束；再不行则退回 PD hover。
"""

import os, time as pytime
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import cvxpy as cp
from math import ceil

# ------------- params -------------
DT = 0.1
T_TOTAL = 60.0
STEPS = int(ceil(T_TOTAL / DT))

# weights
Q_pos = 10.0
Q_vel = 2.0
R_acc = 0.1
Qf_pos = 25.0
Qf_vel = 5.0

# bounds
POS_MIN = np.array([-10.0, -40.0, 0.0])
POS_MAX = np.array([140.0,  40.0, 60.0])
VEL_MAX = 12.0
U_MIN = np.array([-5.0, -5.0, -5.0])
U_MAX = np.array([ 5.0,  5.0,  5.0])

TARGET_POS = np.array([80.0, -20.0, 30.0])
X0 = np.zeros(6); X0[2] = 0.0

HORIZON = 24
PRED_SHOW = 12

POS_TOL = 0.5
VEL_TOL = 0.5
HOVER_TIME = 6.0

WAYPOINT_MAX_SPEED = 6.0

# Obstacles: list of (min, max) arrays [x,y,z] pairs
OBSTACLES = [
    (np.array([35.0, -8.0, 0.0]), np.array([55.0, 8.0, 22.0])),
    (np.array([60.0, -30.0, 0.0]), np.array([70.0, -10.0, 18.0])),
    # 可按需增加
]
SAFETY_MARGIN = 2.0

# candidate sides per obstacle
SIDES = ['left', 'right', 'top']   # left -> y <= box_min_y - margin; right -> y >= box_max_y + margin; top -> z >= box_max_z + margin
# cap on combos to prevent explosion
MAX_COMBOS = 27

SHOW_INTERACTIVE = True
SAVE_FIG = True

# --------- dynamics & costs ----------
I3 = np.eye(3)
A = np.block([[I3, DT * I3],[np.zeros((3,3)), I3]])
B = np.block([[0.5 * DT**2 * I3],[DT * I3]])
Q = np.diag([Q_pos, Q_pos, Q_pos, Q_vel, Q_vel, Q_vel])
Qf = np.diag([Qf_pos, Qf_pos, Qf_pos, Qf_vel, Qf_vel, Qf_vel])
R = R_acc * np.eye(3)

# ---------- refs ----------
def generate_speed_limited_ref(x_start, target_pos, dt, max_speed, steps):
    start_pos = x_start[0:3].copy()
    vec = target_pos - start_pos
    dist = np.linalg.norm(vec)
    if dist < 1e-6:
        return np.tile(target_pos, (steps,1))
    dir_unit = vec / dist
    max_step = max_speed * dt
    n = int(np.ceil((dist - 0.5)/max_step)) if dist > 0.5 else 1
    n = max(1, n)
    seq = [start_pos + dir_unit * min(max_step*i, dist) for i in range(n+1)]
    if len(seq) < steps:
        seq.extend([target_pos] * (steps - len(seq)))
    return np.array(seq[:steps])

ref_pos_seq = generate_speed_limited_ref(X0, TARGET_POS, DT, WAYPOINT_MAX_SPEED, STEPS)
def ref_state_at(i):
    p = ref_pos_seq[min(i, ref_pos_seq.shape[0]-1)]
    xr = np.zeros(6); xr[0:3]=p; xr[3:6]=0.0
    return xr

# ---------- utilities: generate combos ----------
def generate_side_combinations(n_obstacles, sides=SIDES, max_combos=MAX_COMBOS):
    # produce list of tuples length n_obstacles, each element is a side string
    # limit total combos to max_combos by truncation (if n small it's fine)
    import itertools
    all_combos = list(itertools.product(sides, repeat=n_obstacles))
    if len(all_combos) > max_combos:
        # simple heuristic: prefer combos that use 'top' more (safer), then right, left
        def score(combo):
            return sum((0 if s=='top' else (1 if s=='right' else 2)) for s in combo)
        all_combos.sort(key=score)
        all_combos = all_combos[:max_combos]
    return all_combos

# ---------- QP solver that accepts side constraints for obstacles ----------
def solve_mpc_qp_with_sides(x_curr, x_ref, horizon, A, B, Q, R, Qf, u_min, u_max,
                            pos_min=None, pos_max=None, vel_max=None,
                            obstacles=None, combo=None, margin=SAFETY_MARGIN):
    """
    combo: tuple of side strings of length len(obstacles), e.g. ('top','right', 'left')
    Returns: u0, X_pred, U_pred, cost
    Raises runtimeerror if infeasible
    """
    nx = A.shape[0]; nu = B.shape[1]
    X = cp.Variable((nx, horizon+1))
    U = cp.Variable((nu, horizon))
    cost = 0
    cons = []
    cons.append(X[:,0] == x_curr)

    for k in range(horizon):
        cons.append(X[:,k+1] == A @ X[:,k] + B @ U[:,k])
        cons.append(U[:,k] <= u_max); cons.append(U[:,k] >= u_min)
        if pos_min is not None:
            cons.append(X[0:3,k] >= pos_min)
        if pos_max is not None:
            cons.append(X[0:3,k] <= pos_max)
        if vel_max is not None:
            cons.append(cp.norm_inf(X[3:6,k]) <= vel_max)
        # add obstacle-side linear constraints
        if obstacles is not None and combo is not None:
            for idx, side in enumerate(combo):
                bmin, bmax = obstacles[idx]
                if side == 'left':
                    # y <= bmin_y - margin  --> X[1,k] <= val
                    cons.append(X[1,k] <= float(bmin[1] - margin))
                elif side == 'right':
                    cons.append(X[1,k] >= float(bmax[1] + margin))
                elif side == 'top':
                    cons.append(X[2,k] >= float(bmax[2] + margin))
                else:
                    raise ValueError("unknown side")
        # stage cost
        x_err = X[:,k] - x_ref
        cost += cp.quad_form(x_err, Q) + cp.quad_form(U[:,k], R)

    # terminal constraints & cost
    if pos_min is not None:
        cons.append(X[0:3,horizon] >= pos_min)
    if pos_max is not None:
        cons.append(X[0:3,horizon] <= pos_max)
    if vel_max is not None:
        cons.append(cp.norm_inf(X[3:6,horizon]) <= vel_max)
    if obstacles is not None and combo is not None:
        for idx, side in enumerate(combo):
            bmin, bmax = obstacles[idx]
            if side == 'left':
                cons.append(X[1,horizon] <= float(bmin[1] - margin))
            elif side == 'right':
                cons.append(X[1,horizon] >= float(bmax[1] + margin))
            elif side == 'top':
                cons.append(X[2,horizon] >= float(bmax[2] + margin))

    x_errN = X[:,horizon] - x_ref
    cost += cp.quad_form(x_errN, Qf)

    prob = cp.Problem(cp.Minimize(cost), cons)
    prob.solve(solver=cp.OSQP, warm_start=True, eps_abs=1e-3, eps_rel=1e-3, verbose=False, max_iter=5000)
    if prob.status not in [cp.OPTIMAL, cp.OPTIMAL_INACCURATE]:
        raise RuntimeError(f"QP infeasible: {prob.status}")
    u0 = U[:,0].value
    return u0, X.value, U.value, float(prob.value)

# ---------- PD hover ----------
def pd_hover_control(pos, vel, target_pos, kp=0.6, kd=1.0):
    pos_err = pos - target_pos
    return -kp * pos_err - kd * vel

# ---------- plotting helpers ----------
def box_draw(ax, bmin, bmax, color='k', alpha=0.15):
    xs = [bmin[0], bmax[0]]; ys = [bmin[1], bmax[1]]; zs = [bmin[2], bmax[2]]
    for x in xs:
        for y in ys:
            ax.plot([x,x],[y,y],[zs[0],zs[1]], color=color, alpha=0.6)
    for x in xs:
        for z in zs:
            ax.plot([x,x],[ys[0],ys[1]],[z,z], color=color, alpha=0.6)
    for y in ys:
        for z in zs:
            ax.plot([xs[0],xs[1]],[y,y],[z,z], color=color, alpha=0.6)
    Xp = [xs[0], xs[1], xs[1], xs[0], xs[0]]
    Yp = [ys[0], ys[0], ys[1], ys[1], ys[0]]
    Zp = [zs[1]]*5
    ax.plot(Xp, Yp, Zp, color=color, alpha=alpha)

def setup_plot():
    fig = plt.figure(figsize=(14,8))
    gs = fig.add_gridspec(2,3, width_ratios=[2.8,1,1])
    ax3d = fig.add_subplot(gs[:,0], projection='3d')
    ax_pos = fig.add_subplot(gs[0,1]); ax_ctrl = fig.add_subplot(gs[1,1]); ax_u = fig.add_subplot(gs[0,2]); ax_info = fig.add_subplot(gs[1,2]); ax_info.axis('off')
    traj_line, = ax3d.plot([],[],[], lw=2)
    ax3d.scatter([X0[0]],[X0[1]],[X0[2]], c='b', s=80, marker='o', label='start')
    ax3d.scatter([TARGET_POS[0]],[TARGET_POS[1]],[TARGET_POS[2]], c='r', s=120, marker='x', label='target')
    for bmin,bmax in OBSTACLES:
        box_draw(ax3d, bmin, bmax, color='gray', alpha=0.25)
    pred_line, = ax3d.plot([],[],[], ls='--', lw=1.2, color='orange')
    pos_lines = [ax_pos.plot([],[], label='x')[0], ax_pos.plot([],[], label='y')[0], ax_pos.plot([],[], label='z')[0]]; ax_pos.legend(); ax_pos.grid(True)
    u_lines = [ax_ctrl.plot([],[], label='ax')[0], ax_ctrl.plot([],[], label='ay')[0], ax_ctrl.plot([],[], label='az')[0]]; ax_ctrl.legend(); ax_ctrl.grid(True)
    uraw_lines = [ax_u.plot([],[], label='u_raw_x')[0], ax_u.plot([],[], label='u_raw_y')[0], ax_u.plot([],[], label='u_raw_z')[0]]; ax_u.legend(); ax_u.grid(True)
    info_text = ax_info.text(0.01,0.98,"", va='top', family='monospace', fontsize=9)
    return fig, ax3d, ax_pos, ax_ctrl, ax_u, info_text, traj_line, pred_line, pos_lines, u_lines, uraw_lines

def update_plots(fig, ax3d, ax_pos, ax_ctrl, ax_u, info_text, traj_line, pred_line, pos_lines, u_lines, uraw_lines,
                 t_arr, X_arr, U_arr, Uraw_arr, X_pred=None):
    traj_line.set_data(X_arr[:,0], X_arr[:,1]); traj_line.set_3d_properties(X_arr[:,2])
    if X_pred is not None:
        pred_line.set_data(X_pred[0,:], X_pred[1,:]); pred_line.set_3d_properties(X_pred[2,:])
        vel = X_pred[3:6,:]; scale=0.2
        try:
            ax3d.quiver(X_pred[0,:], X_pred[1,:], X_pred[2,:], vel[0,:]*scale, vel[1,:]*scale, vel[2,:]*scale, length=1.0, normalize=False, color='orange', linewidth=1)
        except Exception:
            pass
    if t_arr.size==1:
        t0=float(t_arr[0]);
        for ax in (ax_pos, ax_ctrl, ax_u):
            ax.set_xlim(t0-0.1,t0+0.1)
    else:
        for ax in (ax_pos, ax_ctrl, ax_u): ax.set_xlim(t_arr[0], t_arr[-1])
    for i,ln in enumerate(pos_lines): ln.set_data(t_arr, X_arr[:,i])
    for i,ln in enumerate(u_lines): ln.set_data(t_arr, U_arr[:,i] if U_arr.size>0 else np.array([]))
    for i,ln in enumerate(uraw_lines): ln.set_data(t_arr, Uraw_arr[:,i] if Uraw_arr.size>0 else np.array([]))
    fig.canvas.draw(); fig.canvas.flush_events()

# ------------- main -------------
def main():
    out_dir = os.path.abspath('.')
    log_path = os.path.join(out_dir, 'mpc_multi_obstacle_log.csv')
    fig_path = os.path.join(out_dir, 'mpc_multi_obstacle.png')

    if SHOW_INTERACTIVE:
        plt.ion()
        fig, ax3d, ax_pos, ax_ctrl, ax_u, info_text, traj_line, pred_line, pos_lines, u_lines, uraw_lines = setup_plot()
    else:
        fig=None

    times=[]; X_hist=[]; U_hist=[]; Uraw_hist=[]
    x = X0.copy()
    hover_mode=False; hover_start=None; reached=False
    combos = generate_side_combinations(len(OBSTACLES))
    print(f"Generated {len(combos)} candidate combos (capped at {MAX_COMBOS})")

    sim_start = pytime.time()
    for step in range(STEPS):
        t = step*DT
        xr = ref_state_at(step)
        best_solution = None
        best_cost = None
        best_Xpred = None
        best_Upred = None
        # enumerate combos
        for combo in combos:
            try:
                u0, Xpred, Upred, cost = solve_mpc_qp_with_sides(x, xr, HORIZON, A,B,Q,R,Qf, U_MIN, U_MAX,
                                                                 pos_min=POS_MIN, pos_max=POS_MAX, vel_max=VEL_MAX,
                                                                 obstacles=OBSTACLES, combo=combo, margin=SAFETY_MARGIN)
                # record feasible
                if best_cost is None or cost < best_cost:
                    best_cost = cost
                    best_solution = u0
                    best_Xpred = Xpred
                    best_Upred = Upred
            except Exception:
                continue

        # If none feasible with strict corridors, try relaxed (no pos_min/max)
        if best_solution is None:
            for combo in combos:
                try:
                    u0, Xpred, Upred, cost = solve_mpc_qp_with_sides(x, xr, HORIZON, A,B,Q,R,Qf, U_MIN, U_MAX,
                                                                     pos_min=None, pos_max=None, vel_max=VEL_MAX,
                                                                     obstacles=OBSTACLES, combo=combo, margin=SAFETY_MARGIN)
                    if best_cost is None or cost < best_cost:
                        best_cost = cost
                        best_solution = u0
                        best_Xpred = Xpred
                        best_Upred = Upred
                except Exception:
                    continue

        # final fallback
        if best_solution is None:
            # no QP feasible -> PD hover fallback
            print(f"[t={t:.2f}] No feasible corridor QP -> PD fallback")
            u_opt = pd_hover_control(x[0:3], x[3:6], TARGET_POS)
            Xpred = None
            Uraw = u_opt
        else:
            u_opt = best_solution
            Xpred = best_Xpred
            Uraw = u_opt

        u_apply = np.clip(u_opt, U_MIN, U_MAX)
        x = A.dot(x) + B.dot(u_apply)
        times.append(t); X_hist.append(x.copy()); U_hist.append(u_apply.copy()); Uraw_hist.append(Uraw.copy() if np.all(np.isfinite(Uraw)) else np.zeros(3))

        # arrival
        pos=x[0:3]; vel=x[3:6]
        err_to_target = np.linalg.norm(pos-TARGET_POS)
        if (err_to_target<POS_TOL) and (np.linalg.norm(vel)<VEL_TOL) and (not hover_mode):
            hover_mode=True; hover_start=t; print(f"*** ARRIVED -> hover {HOVER_TIME}s (t={t:.2f}) ***")
        if hover_mode and (t - hover_start >= HOVER_TIME):
            print("*** HOVER complete -> ending sim ***"); reached=True; break

        # update plot
        try:
            if SHOW_INTERACTIVE and fig is not None:
                t_arr=np.array(times); X_arr=np.array(X_hist); U_arr=np.array(U_hist); Uraw_arr=np.array(Uraw_hist)
                update_plots(fig, ax3d, ax_pos, ax_ctrl, ax_u, info_text, traj_line, pred_line, pos_lines, u_lines, uraw_lines,
                             t_arr, X_arr, U_arr, Uraw_arr, X_pred=Xpred)
                plt.pause(0.001)
        except Exception as e:
            print("Plot error:", e); break

    sim_end=pytime.time()
    print(f"Sim done. wall={sim_end-sim_start:.2f}s steps={len(times)} reached={reached}")

    # save logs & fig
    Xnp=np.array(X_hist); Unp=np.array(U_hist); Urawnp=np.array(Uraw_hist)
    df=pd.DataFrame({'t':np.array(times),'x':Xnp[:,0],'y':Xnp[:,1],'z':Xnp[:,2],
                     'vx':Xnp[:,3],'vy':Xnp[:,4],'vz':Xnp[:,5],
                     'ax':Unp[:,0],'ay':Unp[:,1],'az':Unp[:,2],
                     'u_raw_x':Urawnp[:,0],'u_raw_y':Urawnp[:,1],'u_raw_z':Urawnp[:,2]})
    df.to_csv(log_path, index=False)
    print("Saved CSV:", log_path)
    try:
        if SAVE_FIG and fig is not None:
            fig.savefig(fig_path, dpi=160, bbox_inches='tight'); print("Saved figure:", fig_path)
    except Exception as e:
        print("Save fig failed:", e)
    if SHOW_INTERACTIVE:
        plt.ioff(); plt.show()


if __name__=='__main__':
    main()
